import numpy as np
import cv2
import glob
import os

# 1. 相机标定
def calibrate_camera(chessboard_size, calibration_images_path):
    objp = np.zeros((chessboard_size[0] * chessboard_size[1], 3), np.float32)
    objp[:, :2] = np.mgrid[0:chessboard_size[0], 0:chessboard_size[1]].T.reshape(-1, 2)

    obj_points = []  # 3d points in real world space
    img_points = []  # 2d points in image plane

    images = glob.glob(os.path.join(calibration_images_path, '*.jpg'))

    for fname in images:
        img = cv2.imread(fname)
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        ret, corners = cv2.findChessboardCorners(gray, chessboard_size, None)

        if ret:
            obj_points.append(objp)
            img_points.append(corners)

    ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(obj_points, img_points, gray.shape[::-1], None, None)
    return mtx, dist, rvecs, tvecs

# 2. 雷达数据解析
def parse_radar_data(radar_data_path):
    radar_data = np.loadtxt(radar_data_path, delimiter=',')
    return radar_data

# 3. 雷达数据转换
def transform_radar_to_camera(radar_data, rvec, tvec):
    R, _ = cv2.Rodrigues(rvec)
    RT = np.hstack((R, tvec))

    radar_data_homogeneous = np.hstack((radar_data, np.ones((radar_data.shape[0], 1))))
    camera_coords = np.dot(RT, radar_data_homogeneous.T).T

    return camera_coords[:, :3]

# 4. 相机坐标转换为图像坐标系
def camera_to_image_coordinates(camera_coords, camera_matrix):
    points_2d = np.dot(camera_matrix, camera_coords.T).T
    points_2d /= points_2d[:, 2].reshape(-1, 1)  # 转换为齐次坐标
    return points_2d[:, :2]  # 返回(x, y)

# 5. 图像坐标转换为像素坐标
def image_coordinates_to_pixel_coordinates(image_coords, camera_matrix):
    # 提取相机内参数
    f_x = camera_matrix[0, 0]  # x轴焦距
    f_y = camera_matrix[1, 1]  # y轴焦距
    c_x = camera_matrix[0, 2]  # 光心x
    c_y = camera_matrix[1, 2]  # 光心y

    # 转换为像素坐标
    pixel_coords = np.zeros((image_coords.shape[0], 2))
    pixel_coords[:, 0] = (image_coords[:, 0] * f_x) + c_x  # 像素x坐标
    pixel_coords[:, 1] = (image_coords[:, 1] * f_y) + c_y  # 像素y坐标

    return pixel_coords

# 6. 主函数
def main():
    chessboard_size = (9, 6)  # 内部角点的数量
    calibration_images_path = 'calibration_images'  # 存放棋盘格图像的路径
    radar_data_path = 'radar_data.csv'  # 雷达数据的路径

    # 相机标定
    camera_matrix, dist_coeffs, rvecs, tvecs = calibrate_camera(chessboard_size, calibration_images_path)
    print("Camera Matrix:\n", camera_matrix)

    # 假设使用第一个标定得到的旋转向量和位移向量
    rvec = rvecs[0]
    tvec = tvecs[0]

    # 解析雷达数据
    radar_data = parse_radar_data(radar_data_path)

    # 转换雷达数据到相机坐标
    camera_coordinates = transform_radar_to_camera(radar_data, rvec, tvec)
    print("Camera Coordinates:\n", camera_coordinates)

    # 将相机坐标转换为图像坐标系
    image_coordinates = camera_to_image_coordinates(camera_coordinates, camera_matrix)
    print("Image Coordinates:\n", image_coordinates)

    # 将图像坐标转换为像素坐标
    pixel_coordinates = image_coordinates_to_pixel_coordinates(image_coordinates, camera_matrix)
    print("Pixel Coordinates:\n", pixel_coordinates)

# 测试代码
if __name__ == '__main__':
    main()